package top.lingkang.finalgateway.core.ws;

import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.HttpHeaders;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import io.netty.handler.codec.http.websocketx.WebSocketFrame;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okio.ByteString;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import top.lingkang.finalgateway.config.RouterConfig;
import top.lingkang.finalgateway.core.RequestHandler;
import top.lingkang.finalgateway.config.NextTarget;

import java.util.Map;

/**
 * @author lingkang
 * @create by 2024/5/10 15:22
 */
@Slf4j
public class WebSocketInitializer extends SimpleChannelInboundHandler<WebSocketFrame> {

    private HttpHeaders headers;
    private RequestHandler handler;
    private ChannelHandlerContext ctx;
    private WebSocket ws;
    private RouterConfig config;
    private long total = 0;
    private NextTarget next;

    private WebSocketListener listener = new WebSocketListener() {
        @Override
        public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
            if (ctx.channel().isActive()) {
                ctx.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }

        @Override
        public void onClosing(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        }

        @Override
        public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
            // log.warn("onFailure", t);
            if (ctx.channel().isActive()) {
                ctx.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }

        @Override
        public void onMessage(@NotNull WebSocket webSocket, @NotNull ByteString bytes) {
            ctx.writeAndFlush(new BinaryWebSocketFrame(
                    Unpooled.copiedBuffer(bytes.getData$okio())
            ));
        }

        @Override
        public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
            ctx.writeAndFlush(new TextWebSocketFrame(text));
        }

        @Override
        public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
        }
    };

    public WebSocketInitializer(HttpHeaders headers, RequestHandler handler, RouterConfig config) {
        this.headers = headers;
        this.handler = handler;
        this.config = config;
    }

    @Override
    public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
        super.handlerAdded(ctx);
        this.ctx = ctx;

        next = config.nextTarget(handler);
        // open
        Request.Builder builder = new Request.Builder()
                .url(next.getTarget());
        String referer = headers.get("referer");
        if (referer != null) {
            String urlHost = headers.get("host");
            if (urlHost == null) {
                headers.remove("referer");
            } else {
                headers.set("referer", referer.replace(urlHost, next.getHost()));
            }
        }
        headers.remove("Sec-WebSocket-Extensions");
        headers.set("host", next.getHost());
        for (Map.Entry<String, String> e : headers) {
            builder.header(e.getKey(), e.getValue());
        }
        Request request = builder.build();
        ws = config.getClient().newWebSocket(request, listener);
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, WebSocketFrame msg) throws Exception {
        // message
        if (msg instanceof TextWebSocketFrame) {
            String text = ((TextWebSocketFrame) msg).text();
            total += text.length();
            ws.send(text);
        } else if (msg instanceof BinaryWebSocketFrame) {
            BinaryWebSocketFrame bw = (BinaryWebSocketFrame) msg;
            byte[] array = bw.content().array();
            total += array.length;
            ws.send(new ByteString(array));
        } else {
            log.warn("Unknown WebSocketFrame type: " + msg.getClass().getName());
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        // exception
        log.warn("exceptionCaught", cause);
    }

    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
        if (ws != null) {
            ws.cancel();
            // ws.close(1000, "close");
        }
        config.getMonitor().accessWs(ctx.channel().remoteAddress().toString(), handler.getUri(), next.getTarget(), total);
    }
}
